#ifndef TESTSTREEVIEW_H
#define TESTSTREEVIEW_H

#include <atlbase.h>
#include <atlctrls.h>
#include <atlsimpcoll.h>
#include <atlcomcli.h> // for CComBSTR, for GetText
#include <atlutil.h>
#include <atltypes.h>

#include "UiInterfaces.h"
#include "TestRunner.h"
#include "resource.h"

#include <map>
#include <vector>
#include <algorithm>
using namespace std;

namespace TddGui
{

#define IS_GREY_STATE(x) (x)&0x000 // including Grey, GreenGrey, RedGrey
#define IS_YELLOW_STATE(x) (x)&0x0100 // including Yellow, GreenYellow, RedYellow
#define OPPSITE_STATE(x) (x)^0x0111

enum TestNodeStates
{
	Grey = 0x000, // not to run
	GreenGrey = 0x001,// last pass, but not to run
	RedGrey = 0x010,  // last fail, but not to run
	Yellow = 0x111,    // to be run test
	GreenYellow = 0x110, // last pass, but to be run test
	RedYellow = 0x101 // last fail, but to be run test
};
map<TestNodeStates, int> colorIndexMap;
map<int, TestNodeStates> indexColorMap;

template<typename TestRunner>
class TestsTreeViewT : public CWindowImpl<TestsTreeViewT<TestRunner>, CTreeViewCtrlEx>, public LoadNotification, public RunNotification
{
	EnableNotification& m_run;
	TextNotification& m_staticText;
	TextNotification& m_errorText;
	ProgressNotification& m_progress;

	struct ErrorArray : private CSimpleArray<BSTR>
	{
		~ErrorArray() { RemoveAll(); }
		void RemoveAll()
		{
			for (int i=0; i<GetSize(); ++i)
			{
				SysFreeString(this->operator[](i));
			}
			CSimpleArray<BSTR>::RemoveAll();
		}
		using CSimpleArray<BSTR>::Add;
		using CSimpleArray<BSTR>::operator[];
		using CSimpleArray<BSTR>::GetSize;
	} m_errors;
	CPath m_pathDll;

public:
	TestsTreeViewT (TextNotification& subjectStaticText,
					TextNotification& subjectErrorText,
					EnableNotification& subjectEnable,
					ProgressNotification& subjectProgress)
		: m_run       (subjectEnable)
		, m_staticText(subjectStaticText)
		, m_errorText (subjectErrorText)
		, m_progress  (subjectProgress)
	{
		colorIndexMap[Grey] = 0;
		indexColorMap[0] = Grey;

		colorIndexMap[GreenGrey] = 1;
		indexColorMap[1] = GreenGrey;

		colorIndexMap[RedGrey] = 2;
		indexColorMap[2] = RedGrey;

		colorIndexMap[Yellow] = 3;
		indexColorMap[3] = Yellow;

		colorIndexMap[GreenYellow] = 4;
		indexColorMap[4] = GreenYellow;

		colorIndexMap[RedYellow] = 5;
		indexColorMap[5] = RedYellow;
	}

	virtual ~TestsTreeViewT() {}
	
	void Init(HWND hwnd)
	{
		SubclassWindow(hwnd);

		CImageList il;
		il.Create(IDB_ERRORS, 16, 3, RGB(0,0,0));
		SetImageList(il, TVSIL_NORMAL);
	}

	BEGIN_MSG_MAP(TestsTreeViewT)
		REFLECTED_NOTIFY_CODE_HANDLER(NM_DBLCLK, OnDBClick)
		//ON_NOTIFY(NM_CLICK, IDC_TREE_TESTS, OnClick);
		MESSAGE_HANDLER(WM_DESTROY, OnDestroy)
        REFLECTED_NOTIFY_CODE_HANDLER(TVN_SELCHANGED, OnItemSelected)
        DEFAULT_REFLECTION_HANDLER();
    END_MSG_MAP()

private:
    LRESULT OnDestroy(UINT, WPARAM, LPARAM, BOOL& bHandled)
	{
		SetImageList(NULL).Destroy();
		return bHandled = TRUE;
	}
	LRESULT OnItemSelected (int, LPNMHDR, BOOL& bHandled)
	{
		SetErrorTextIfAny();
		return bHandled = TRUE;
	}

	// the design for the following 3 traverse is so bad, and need to refactor
	// traverse downwards to log disable nodes' name
	void DownTreeLogDisableName(CTreeItem root, vector<CString>& result)
	{	
		// check whether this node is disabled
		int imageIndex, selectedImageIndex;
		root.GetImage(imageIndex, selectedImageIndex);

		TestNodeStates color = indexColorMap[imageIndex];
		bool bIsYellow = IS_YELLOW_STATE(color);
		if ( !bIsYellow )
		{
			// it is disabled, log the name down
			CString fullTestName;
			CTreeItem node = root;
			while(!node.GetParent().IsNull())
			{
				CComBSTR bstr;
				node.GetText(bstr.m_str);
				if (!fullTestName.IsEmpty())
					fullTestName = L"::" + fullTestName;
				fullTestName = bstr + fullTestName;

				node = node.GetParent();
			}

			result.push_back(fullTestName);
		}

		// it is not the leaf, set all the sub-tree nodes as the opposite of the root
		CTreeItem temp = root.GetChild();
		while(!temp.IsNull())
		{
			DownTreeLogDisableName(temp, result );
			temp = temp.GetNextSibling();
		}
	}

	// traverse downwards to set disable nodes' name
	// this function is called when the dll is newly loaded, and we want to resotre the disable modes
	void DownTreeSetDisableState(CTreeItem root, CString prefixName, vector<CString>& disableUtFullName)
	{	
		// check whether this node is disabled
		// first get the node's full name
		CComBSTR bstr;
		root.GetText(bstr.m_str);
		CString nodeName = prefixName + bstr;
		
		//check is it in?
		if (disableUtFullName.end() != find(disableUtFullName.begin(), disableUtFullName.end(), nodeName))
		{
			// disable this node
			// it must to make the nodes grey, for the initial state is yellow
			ChangeColor(root);
		}

		// it is not the leaf, set all the sub-tree nodes as the opposite of the root
		CTreeItem temp = root.GetChild();
		while(!temp.IsNull())
		{
			prefixName = nodeName + L"::";
			DownTreeSetDisableState(temp, prefixName, disableUtFullName);
			temp = temp.GetNextSibling();
		}
	}

	// traverse downwards to change the node's color
	void DownTreeChangeColor(CTreeItem root, TestNodeStates color)
	{	
		ChangeColor(root, &color);
		
		// it is not the leaf, set all the sub-tree nodes as the opposite of the root
		CTreeItem temp = root.GetChild();
		while(!temp.IsNull())
		{
			DownTreeChangeColor(temp, color );
			temp = temp.GetNextSibling();
		}
	}
	
	void ChangeParentEnable(CTreeItem item)
	{
		CTreeItem parent = item.GetParent();
		while(!parent.IsNull())
		{
			int originIndex, changedIndex;
			parent.GetImage(originIndex, changedIndex);

			// if the item is disable, enable it
			if (!(IS_YELLOW_STATE(indexColorMap[originIndex])) )
			{
				// grey state, change to the yellow one
				TestNodeStates image = (TestNodeStates)(OPPSITE_STATE(indexColorMap[originIndex]));
				parent.SetImage(colorIndexMap[image], colorIndexMap[image]);
			}

			parent = parent.GetParent();
		}
	}

	// return the state changes
	TestNodeStates ChangeColor(CTreeItem hItem, TestNodeStates *pColor = NULL)
	{
		int imageIndex, selectedImageIndex;

		// get the original color
		hItem.GetImage(imageIndex, selectedImageIndex);

		if (pColor == NULL)
		{
			// change to the opposite color
			TestNodeStates temp = (TestNodeStates)(OPPSITE_STATE(indexColorMap[imageIndex]));
			imageIndex = colorIndexMap[temp];
		}
		else
		{	
			// change the special color
			imageIndex = colorIndexMap[*pColor];
		}

		hItem.SetImage(imageIndex, imageIndex);	
		
		return indexColorMap[imageIndex];
	}

	LRESULT OnDBClick(int, LPNMHDR, BOOL &bHandled)
	{
		DWORD dwPos = GetMessagePos(); 
		CPoint point(LOWORD(dwPos), HIWORD(dwPos) ); 
		ScreenToClient(&point); 

		UINT uFlag; 
		HTREEITEM hItem = HitTest(point, &uFlag); 
		if(uFlag & TVHT_ONITEM) 
		{
			CTreeItem selected(hItem, this);
			
			// test which color it will use
			TestNodeStates changedColor = ChangeColor(selected);

			if (selected.HasChildren())
			{
				// select the root
				DownTreeChangeColor(selected, changedColor);
			}

			// does it become enable?
			if (IS_YELLOW_STATE(changedColor)) // all Yellow begins with 1, all Gray begins with 0
			{
				// change the selected node's parent's color
				ChangeParentEnable(selected);
			}
		}

		return bHandled = false;
	}

	void SetErrorTextIfAny()
	{
		LPCWSTR errorText = L"";

		CTreeItem selected = GetSelectedItem();
		if (!selected.IsNull())
		{
			// the error msg is stored as the data of the tree node
			DWORD_PTR lparam = selected.GetData();
			if (lparam != NULL)
				errorText = reinterpret_cast<LPCWSTR>(lparam);
		}
		m_errorText.SetText(errorText);
	}
	virtual void LoadTests(const CPath& path)
	{
		HRESULT hr = TestRunner::IsTestDll(path);
		if (SUCCEEDED(hr)) {
			// form the tree view
			PopulateTreeView(path);

			SelectItem(GetRootItem());
			m_staticText.SetText(path);
			m_run.Enable(true);
		} else if (hr == E_FAIL) {
			MessageBox(L"no tests found in " + CString(path), L"TddGui");
		} else {
			MessageBox(L"error loading " + CString(path) + L": " + AtlGetErrorDescription(hr), L"TddGui");
		}
	}

	void PopulateTreeView(const CPath& path)
	{
		m_pathDll = path;

		DeleteAllItems();
		m_errors.RemoveAll();

		CTreeItem root = InsertItem(L"All Tests", TVI_ROOT, TVI_LAST);
		root.SetImage(colorIndexMap[Yellow], colorIndexMap[Yellow]);

		class InsertNodePredicate : public TDD::Discriminator
		{
			CTreeItem& m_root;
		public:
			InsertNodePredicate(CTreeItem& root) : m_root(root) { }
			virtual bool WantTest(const TDD::UnitTestInfo& uti)
			{
				CString fullName(uti.group);
				fullName += L"::";
				fullName += CString(uti.testname);

				// format of the fullName is either
				// Outer::Inner::Class::Method, or
				// Class::Method

				CTreeItem node = m_root;

				int iStart = 0;
				for(;;)
				{
					CString cs = fullName.Tokenize(L"::", iStart);
					if (iStart == -1)
						break;

					// to create the branch
					// for each namespace, class or method, see if it's in there already and add it if not
					CTreeItem child = node.GetChild();
					if (child.IsNull())
					{
						node = node.AddTail(cs, colorIndexMap[Yellow]);
						continue;
					}

					for(;;)
					{
						// first traverse the tree branch to the test belongs to 
						CComBSTR bstrText;
						child.GetText(bstrText.m_str);
						if (cs == bstrText)
						{
							node = child;
							break;	
						}

						// now we have found the tree branch, and can add the leaves
						CTreeItem sibling = child.GetNextSibling();
						if (sibling.IsNull())
						{
							node = child.GetParent().AddTail(cs, colorIndexMap[Yellow]);
							break;
						}
						child = sibling;
					}
				}
				return false;
			}
		};

		InsertNodePredicate inp(root);
		TestRunner::RunTddTestsInDll(path, inp, TestRunner::GetNullReporter());
	}

	void SetState(vector<CString>& disableNames)
	{

	}

private:
	virtual void Run(void)
	{
		// from the one we have selected to begin testing
		CTreeItem selected = GetSelectedItem();
		if (selected.IsNull())
		{
			MessageBox(L"no test(s) selected!", L"TddGui");
			return;
		}

		// first lets log the disabled unit tests, we must log their full name
		vector<CString> disableNodeNames;
		// the disable node name will not contain "All Tests"
		DownTreeLogDisableName(GetRootItem(), disableNodeNames); 

		// reset progress bar
		m_progress.Status(0, 100, RGB(0xFF, 0xFF, 0xFF));

		// construct (full) test name, class or namespace, but not including "All Tests"
		CString fullTestName;
		CTreeItem node = selected;
		while(!node.GetParent().IsNull())
		{
			CComBSTR bstr;
			node.GetText(bstr.m_str);
			if (!fullTestName.IsEmpty())
				fullTestName = L"::" + fullTestName;
			fullTestName = bstr + fullTestName;

			node = node.GetParent();
		}

		// in case the .dll got rebuilt
		PopulateTreeView(m_pathDll);

		// set the state again according to the name vector if the tests still exist
		DownTreeSetDisableState(GetRootItem().GetChild(), L"", disableNodeNames);

		// try find the selected item again (it might not be there any more, if the .dll got rebuilt)...
		node = FindNode(GetRootItem(), fullTestName);
		if (node.IsNull())
		{
			MessageBox(L"can no longer find selected test(s)", L"TddGui");
			return;
		}

		// ... and reselect it
		node.Select();

		// ok, we're ready.  Run the tests.
		class RunSelectedTests : public TDD::Discriminator
		{
			CString m_fullTestNodeName;
			vector<CString> m_disableNodeNames;
		public:
			RunSelectedTests(
				const CString& fullTestName, 
				vector<CString>& disableNodeNames)
				: m_fullTestNodeName(fullTestName), 
				  m_disableNodeNames(disableNodeNames)
			{}

			virtual bool WantTest(const TDD::UnitTestInfo& uti)
			{
				bool bResult = false;
				CString fullName(uti.group);
				// check whether the node is in the sub-tree of the selected node
				fullName += L"::";
				fullName += CString(uti.testname);
				if (0 != fullName.Find(m_fullTestNodeName)) // must match starting at position 0
					return false;

				// we can not rely on "All Tests" any more
				if (m_fullTestNodeName.IsEmpty())
				{
					// the selected node is the top node, all other nodes are his children
					bResult = true;
				}
				else
				{
					// but it's still possible that this isn't a match.  Example:
					// Slow::ExpressionsAndOperandsTests::MsiAttributes - the selected test, but the code above also matches
					// Slow::ExpressionsAndOperandsTests::MsiAttributesUsingDefault.
					// idea:  we can check the next char after m_fullTestNodeName; if it's a ':' or '\0' we've got a match
					// this works because we need to be either at the end of a namespace or class boundary (:), or else
					// a full single test name (\0).
					WCHAR wc = fullName[m_fullTestNodeName.GetLength()];
					if (wc == L':' || wc == L'\0') // namespace/class boundary, or single selected test
					{
						bResult = true; 
					}
					else
					{
						bResult = false;
					}
				}
				
				// it is in the selected node, now check whether it is disabled
				if (bResult == true) /*name is ok*/
				{
					if (m_disableNodeNames.end() != 
						find(m_disableNodeNames.begin(), m_disableNodeNames.end(), fullName))
					{
						// it is disabled node
						bResult = false;
					}
				}
				
				return bResult;
			}
		};

		// not so fast:  first, count up how many tests we're going to run so that the progress bar works
		class CountSelectedTests : public RunSelectedTests
		{
			unsigned int& m_count;
		public:
			CountSelectedTests(const CString& fullTestName, unsigned int& count, vector<CString>& disableNames)
				: RunSelectedTests(fullTestName, disableNames)
				, m_count(count)
			{}
		private:
			virtual bool WantTest(const TDD::UnitTestInfo& uti)
			{
				// here to check how many unit tests are ready to run
				if (RunSelectedTests::WantTest(uti) == true)
					++m_count;
				return false;
			}
		};
		
		// attention here, we need to count up how many tests we are going to run
		unsigned int count = 0;
		CountSelectedTests cst(fullTestName, count, disableNodeNames);
		TestRunner::RunTddTestsInDll(m_pathDll, cst, TestRunner::GetNullReporter());

		class Status
		{
			unsigned int m_current, m_total;
			COLORREF m_color;
			ProgressNotification& m_progress;
		public:
			Status(unsigned int total, ProgressNotification& subjectProgress)
				: m_current(0)
				, m_total(total)
				, m_color(RGB(0x00, 0xFF, 0x00))
				, m_progress(subjectProgress)
			{}
			void OnNextTest()
			{
				++m_current;
				SendStatus();
			}
			void OnError()
			{
				m_color = RGB(0xFF, 0x00, 0x00);
				SendStatus();
			}
		private:
			void SendStatus() { m_progress.Status(m_current, m_total, m_color); }
		};
		Status status(count, m_progress);

		class RunTestsWithProgress : public RunSelectedTests
		{
			Status& m_status;
		public:
			RunTestsWithProgress(const CString& fullTestName, Status& status, vector<CString>& disableNames)
				: RunSelectedTests(fullTestName, disableNames)
				, m_status(status)
			{}
		private:
			virtual bool WantTest(const TDD::UnitTestInfo& uti)
			{
				bool wantTest = RunSelectedTests::WantTest(uti);
				if (wantTest)
					m_status.OnNextTest();
				return wantTest;
			}
		};

		class Colorizer : public TDD::Reporter
		{
			CTreeItem m_root, m_current;
			ErrorArray& m_errors;
			Status& m_status;
		public:
			Colorizer(Status& status, ErrorArray& errors, CTreeItem root)
				: m_errors(errors)
				, m_root(root)
				, m_status(status)
			{}
		private:
			virtual void ForEachTest   (const TDD::UnitTestInfo& uti)
			{
				CString fullName(uti.group);
				fullName += L"::";
				fullName += CString(uti.testname);

				m_current = TestsTreeViewT::FindNode(m_root, fullName);
				// this time runnable, next time must be runnable
				m_current.SetImage(colorIndexMap[GreenYellow], colorIndexMap[GreenYellow]);

				CTreeItem node = m_current.GetParent();
				while(!node.IsNull())
				{
					int imageIndex, selectedIndex;
					node.GetImage(imageIndex, selectedIndex);
					
					if (indexColorMap[imageIndex] != Yellow)
						break;
					node.SetImage(colorIndexMap[GreenYellow], colorIndexMap[GreenYellow]);
					node = node.GetParent();
				}
			}
			virtual void ForEachFailure(const TDD::TestFailure& failure)
			{
				m_status.OnError();

				// this times runnable, next time must be runnable
				m_current.SetImage(colorIndexMap[RedYellow], colorIndexMap[RedYellow]);
				m_current.Select();

				CTreeItem node = m_current.GetParent();
				while(!node.IsNull())
				{
					node.SetImage(colorIndexMap[RedYellow], colorIndexMap[RedYellow]);
					node = node.GetParent();
				}

				CString cs(failure.file_name);
				if (cs.GetLength() > 52)
					cs = L"..." + cs.Right(52);
				CString csLine;
				csLine.Format(L"\r\non line %d:\r\n", failure.line_number);
				cs += csLine;
				cs += failure.error_string;
				
				m_errors.Add(SysAllocString(cs));
				m_current.SetData(reinterpret_cast<DWORD_PTR>(m_errors[m_errors.GetSize()-1]));
			}
		};
 
		RunTestsWithProgress rtwp(fullTestName, status, disableNodeNames);
		Colorizer colorizer(status, m_errors, GetRootItem());
		TestRunner::RunTddTestsInDll(m_pathDll, rtwp, colorizer);
		SetErrorTextIfAny();
	}

private:
	static CTreeItem FindNode(const CTreeItem& root, const CString& fullTestName)
	{
		CTreeItem node = root;

		int iStart = 0;
		for(;;)
		{
			CString cs = fullTestName.Tokenize(L"::", iStart);
			if (iStart == -1)
				break;

			CTreeItem child = node.GetChild();
			if (child.IsNull())
			{
				node = child; // not found; can test for null-ness, below
				break;
			}

			for(;;)
			{
				CComBSTR bstrText;
				child.GetText(bstrText.m_str);
				if (cs == bstrText)
				{
					node = child;
					break;	
				}
				CTreeItem sibling = child.GetNextSibling();
				if (sibling.IsNull())
				{
					node = sibling; // not found; can test for null-ness, below
					break;
				}
				child = sibling;
			}
			if (node.IsNull())
				break;
		}
		return node;
	}
};
typedef TestsTreeViewT<TestRunner> TestsTreeView;
}

#endif